import numpy as np
import sys
sys.path.append("/Users/ender/Code_repo/rosalind-brush-questions")
from utils import read_fasta


fasta = read_fasta("../10_ConsensusAndProfile/rosalind_cons.txt")

seq_length = len(fasta[list(fasta.keys())[0]])
rownames = np.array(list("ACGT"))

res = np.zeros((4, seq_length), dtype=int)
for i in range(seq_length):
    for name, seq in fasta.items():
        if seq[i] == "A":
            res[0, i] += 1
        elif seq[i] == "C":
            res[1, i] += 1
        elif seq[i] == "G":
            res[2, i] += 1
        elif seq[i] == "T":
            res[3, i] += 1

# Output
print("".join(rownames[np.argmax(res, axis=0)]))
res = res.astype(str)
for i in range(4):
    print(f"{rownames[i]}:", " ".join(res[i, :]))
